簡介

Text Generation Inference 簡稱 TGI，是由 Hugging Face 開發的 LLM Inference 框架。其中整合了相當多推論技術，例如 Flash Attention, Paged Attention, Continuous Batching 以及 BNB & GPTQ Quantization 等等，加上 Hugging Face 團隊強大的開發能量與活躍的社群參與，使 TGI 成為部署 LLM Service 的最佳選擇之一。

今天就來介紹 TGI 的用法吧！

架設服務

TGI 提供 Local Install 跟 Docker Image 兩種用法，但是 Local Install 要花上非常久的時間編譯與安裝，而且也很容易遇到環境配置的問題，除非要客製化停用某些套件之類的，不然一律推薦使用 Docker Image 操作。

可以從官方 GitHub Packages 查看 TGI Docker Image 有哪些版本，或者直接使用 latest 最新版本：

docker pull ghcr.io/huggingface/text-generation-inference:latest

筆者目前使用的是 sha-5ba53d4 版，因為這個 Docker Image 大小約 10GB 左右，所以下載需要花一段時間。完成下載之後可以透過 --help 查看參數說明：

docker run --rm ghcr.io/huggingface/text-generation-inference --help

我們可以先用個小模型 OPT-125M 簡單測試一下：

docker run --gpus all --shm-size 1g \
    -p 8080:80 -v $PWD/data:/data \
    ghcr.io/huggingface/text-generation-inference \
    --model-id facebook/opt-125m

等到訊息紀錄出現 Connected 就代表服務啟動完成，我們可以透過 curl 測試：

curl -X POST 127.0.0.1:8080/generate \
    -d '{"inputs":"Hello, "}' \
    -H 'Content-Type: application/json'

得到類似以下的輸出：

{"generated_text":" I'm a new player and I'm looking for a good team to play with. I'm a"}

TGI 會透過網路下載模型，並且把模型轉換成 Safetensors 格式，然後將模型權重存在容器內部的 /data 路徑裡面。可以使用參數 -v $PWD/data:/data 將此路徑 Mapping 出來，這樣就不用每次執行的時候都要重新下載一遍模型權重。反過來說，我們也可以先將模型下載好，然後再 Mapping 進容器裡面，例如：

# launch_tgi.sh

MODEL_PATH=Models/opt-125m

git clone https://huggingface.co/facebook/opt-125m $MODEL_PATH

docker run --gpus all --shm-size 1g \
    -p 8080:80 -v $PWD/$MODEL_PATH:/$MODEL_PATH \
    ghcr.io/huggingface/text-generation-inference \
    --model-id /$MODEL_PATH

基本用法

TGI 的 API 用法可以參考說明文件，這裡介紹幾個比較重要的參數：

將 decoder_input_details 和 details 設為 True 會回傳一些詳細資訊，例如生成 Token 總數、輸入 Prompt 的 Token 總數等等。給 stop 一個 List of String 可以控制輸出停止點。

最特別的是 TGI 可以設定 truncate 參數，讓系統幫你把太長的 Prompt 截斷，例如設定為 500 的話，就會把 Prompt 前面超過 500 個 Tokens 的內容全部切除，只剩最後面 500 Tokens 當輸入。

這在限制輸入長度很好用，例如做 Retrieval-Based Few-Shot Prompting 時，可以取非常多結果出來，把相似度較低的擺在前面，藉由 truncate 參數自然切掉相似度太低的範例，這樣就能在維持輸入長度不會超出系統限制的同時，盡可能的加上更多範例在 Few-Shot Prompt 裡面。

最後透過 Python 呼叫模型生成的程式碼大致如下：

import json
import requests


url = "http://localhost:8080/generate"
params = {
    "inputs": "This is a long prompt maybe, ",
    "parameters": {
        "best_of": 1,
        "details": True,
        "return_full_text": True,
        "decoder_input_details": True,
        "truncate": 4,  # 只保留最後四個 Tokens
        "max_new_tokens": 128,
        "stop": ["\n", "."],
        "do_sample": True,
        "temperature": 0.5,
        "top_k": 10,
        "top_p": 0.95,
    },
}

resp = requests.post(url, json=params)
result = json.loads(resp.text)
print(result)

因為 return_full_text 被設為 True 所以回傳結果會包含 Prompt 與 Generation，雖然結果裡面看起來 Prompt 並沒有被截斷，但是仔細觀察 details 裡面的 prefill 就可以看到實際的 Prompt 其實是從 "maybe" 開始的。

量化選項

在眾多 TGI 參數裡面，最值得注目的就是我們多了很多量化模型的選擇，透過 --quantize <QUANTIZE> 來指定要使用哪一種，目前 TGI 支援 AWQ, EETQ, BNB, GPTQ 等，以下簡單介紹這些參數選項：

• bitsandbytes 8-Bit 量化，雖然速度偏慢，但還是支援最廣泛、穩定的選擇。
• bitsandbytes-nf4 4-Bit 量化，大部分的模型都可以直接使用此選項，資料型態為 BNB-NF4，可能更符合模型權重分佈的一種資料型態。
• bitsandbytes-fp4 4-Bit 量化，與 BNB-NF4 類似，但使用標準的 4-Bit 浮點數資料型態。
• gptq 4-Bit 量化，需要使用做過 GPTQ Post Training 的模型，可以到 HF Hub 上搜尋，例如 TheBloke 提供的 GPTQ 模型。
• awq 4-Bit 量化，類似 GPTQ 需要提供指定格式的模型，也可以參考 TheBloke 提供的 AWQ 模型。
• eetq 8-Bit 量化，應該可以直接用，但這個選項滿新的，感覺還有些 Bug 的樣子。雖然官方說準備棄用 BNB，但現階段來說 BNB 還是相對穩定一些。

有了量化，我們單顯卡平民就天下無敵了！先拿個 Taiwan Llama 13B 熱熱身：

docker run --gpus all --shm-size 1g \
    -p 8080:80 -v $PWD/data:/data \
    ghcr.io/huggingface/text-generation-inference \
    --model-id yentinglin/Taiwan-LLaMa-v1.0 \
    --quantize bitsandbytes

順利跑起來之後，簡單拿個中文測試：

curl -X POST 127.0.0.1:8080/generate \
    -d '{"inputs":"### USER: 嗨 ### ASSISTANT: "}' \
    -H 'Content-Type: application/json'
# {"generated_text":"你好！我今天可以如何協助你？"}

好欸，看到了我們熟悉的母語！

因為 TGI 參數很多很複雜，於是筆者將常用的參數整理成一份 Python Script 放在此 GitHub Gist 上方便使用，此腳本僅供參考，請以自身運行的環境與 TGI 版本為主。

除了 Decoder 以外，像是 T5 或 M2M 之類的 Encoder-Decoder 也可以用透過 TGI 來運行。但如果不是 TGI 特別支援的模型，可能會回到原本 HF Transformers 的實做。一些像是 Tensor-Parallel 或 Flash Attention 的功能就無法用到，但依然有 Continuous Batching 或 Streaming Outputs 這些功能可以用。因此若要使用 Seq2Seq 模型也是可以考慮用 TGI 來部署。

Token 數量設定

若要更精細的部署模型，與 Token 相關的參數至關重要。對於小顯卡而言，可以減少記憶體消耗，對於大顯卡而言，則能充分利用記憶體空間。最主要的三個參數如下：

• --max-input-length 單筆最大輸入長度。
• --max-total-tokens 單筆最大總長度。
• --max-batch-prefill-tokens 所有 Batch 加起來的最大輸入長度。

其中 --max-batch-prefill-tokens 是影響記憶體消耗最關鍵的參數，會影響整個 Prefill 階段能夠容納多少 Token 當輸入。

在 Decoder-Only LM 裡面，通常會使用 Prefill 代表一開始的輸入，而 Generate 則代表後續 Autoregressive Decoding 的過程。一般而言 Transformers 會在 Prefill 階段消耗掉大量記憶體，而後續 Generate 的記憶體消耗增長速度則相對較緩。也就是說如果一次輸入一整篇長文到 LLM 裡面，可能會直接發生記憶體不足的錯誤。但如果是讓 LLM 慢慢生成一篇長文，則可以生成非常長的文章。

決定這些參數值的方向大概分成輸入長度、輸出長度與批次大小。在不同應用情境下，這些設定都會不太一樣。以 Llama 2 13B 為例，模型是以 4K Context Window 進行訓練的，那我們可以分配 3K 給輸入，剩下 1K 給輸出，因此得到：

1. --max-input-length 為 3000。
2. --max-total-tokens 為 3000 + 1000 = 4000。

註：一般而言 Context Length 的 2K, 4K 通常是以 1024 為單位，所以 3K/1K 的分配實際上會是 3 * 1024 = 3072 跟 1024 個 Tokens。

若我們預計同時推論 4 筆輸入，則可以設定參數：

1. --max-batch-prefill-tokens 為 3000 * 4 = 12000

預設 --max-concurrent-requests 為 128 筆，在這樣的設定下，如果使用者每筆輸入都是 3000 個 Tokens，那 TGI 最多只會同時推論 4 個請求，剩下請求的都會被 Queue 起來，等待任何一個生成結束之後再拿出來放進 Batch 裡面繼續生成。

另外可以將 --max-best-of 設為 1 就好，因為我們通常只需要生成一筆結果。最後使用 bitsandbytes-nf4 做 4-Bit 量化，整個指令看起來會像這樣：

docker run --gpus all --shm-size 1g \
    -p 8080:80 -v $PWD/data:/data \
    ghcr.io/huggingface/text-generation-inference:latest \
    --model-id TheBloke/Llama-2-13B-Chat-fp16 \
    --quantize bitsandbytes-nf4 \
    --max-best-of 1 \
    --max-concurrent-requests 128 \
    --max-input-length 3000 \
    --max-total-tokens 4000 \
    --max-batch-prefill-tokens 12000

不過 TGI 其實是動態決定當下的 Batch Size，例如使用者每筆輸入都減為 1500 個 Tokens 的話，那 Batch Size 可能就會增加到同時推論 8 筆。因此只要在不會超出記憶體的情況下，參數 --max-batch-prefill-tokens 能開多大就開多大，確保服務能夠同時處理的 Batch Size 為最大，便能提昇整個生成的吞吐量。

至於具體到底要調到多少才不會 OOM，那就只能慢慢測試了。筆者習慣以 2K 為單位往上遞增做測試，如果 8K 不會爆就試 10K，依此類推。

TGI Client

TGI Client 是用來跟 TGI Server 交流的客戶端介面，因為在 TGI 的文件沒有太多介紹，所以相當容易被人遺忘。首先透過 pip 安裝 text-generation 套件：

pip install text-generation

基本使用方法如下：

from text_generation import Client

client = Client("http://127.0.0.1:8080", timeout=600)

resp = client.generate(
    "Hello",
    max_new_tokens=16,
    stop_sequences=["."],
    do_sample=False,
    truncate=2048,
)

print(resp.generated_text)

其實就是將 requests 的用法做個包裝，並且幫 Response 物件定義了明確的類別，因此輸入 resp. 的時候，會跳出 generated_text 的提示，開發的時候會更方便一點。透過 TGI Client 進行串流輸出也比較方便：

resp = client.generate_stream(
    "Hello",
    max_new_tokens=16,
    stop_sequences=["."],
    do_sample=False,
    truncate=2048,
)

for chunk in resp:
    print(end=chunk.generated_text, flush=True)
print()

更多詳細的用法，可以參考 PyPI 的介紹頁面。

速度測試

使用空 Prompt 生成 128 Tokens，來比較一下 HF Transformers, vLLM 與 TGI 的速度：

HF   FP16 - 1.23 ms
vLLM FP16 - 0.48 ms
TGI  FP16 - 0.45 ms

TGI 的速度與 vLLM 差不多，但 TGI 的優勢在於量化支援較廣泛，因此來比較一下各種量化方式的速度：

FP16       - 0.45 ms
BNB  8-Bit - 0.64 ms
BNB  4-Bit - 0.73 ms
GPTQ 4-Bit - 0.80 ms
AWQ  4-Bit - 0.55 ms

有量化的模型雖然會略慢一些，但基本上都還是比 HF Transformers 快。這裡筆者沒有測試 EETQ 因為跑起來感覺不太正常，等這個量化方式穩定一點再來測試看看。

最後比較各框架單筆推論的速度：

HF   FP16 - 27 ms
ggml FP16 - 20 ms
TGI  FP16 - 18 ms
vLLM FP16 - 18 ms

其實只看一筆的話，大家的速度都還是差不多的。因此如果是那種放在個人電腦裡面，一人一個 LLM 的用途，ggml 的輕量化部署依然是個滿理想的選擇。

翻譯應用

筆者曾經使用 TGI + Taiwan Llama 替人翻譯過一篇約 13000 Tokens 的英文文章，以下分享這個應用的實做。

首先我先將文章存在 content.txt 裡面，並手動分段，約兩三行的內容就放兩個換行。因為文章內容沒有很長，所以手動分段還算可以，主要是為了確保邊界正確。但如果應用在更長的文章上，可能需要考慮一些自動尋找 Boundary 的工具協助了。

接下來用 TGI 將 Taiwan Llama 架起來，使用 BNB-NF4 量化，接著撰寫以下程式碼進行翻譯：

import json
from concurrent.futures import ThreadPoolExecutor

import requests
from tqdm import tqdm


def main():
    # 讀取文章並以 "\n\n" 切成多個 Chunks
    with open("content.txt", "rt", encoding="UTF-8") as fp:
        text = fp.read().strip()
        text = text.split("\n\n")

    # Taiwan Llama 提供的 Prompt Template
    template = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: 請將以下句子從英文翻譯成中文: {} ASSISTANT:"

    # 定義每段 Chunk 的翻譯函式
    def translate(source):
        prompt = template.format(source)
        target = generate(prompt)
        return {"Original": source, "Translate": target}

    # 最多同時發送 128 個 Requests
    results = list()
    with ThreadPoolExecutor(max_workers=128) as executor:
        with tqdm(total=len(text), ncols=80) as progress:
            for res in executor.map(translate, text):
                results.append(res)
                progress.update()

    # 將結果存成 JSON 檔
    with open("results.json", "wt", encoding="UTF-8") as fp:
        json.dump(results, fp, ensure_ascii=False, indent=4)


# 定義發送 HTTP Request 的函式
def generate(prompt):
    url = "http://localhost:8080/"

    # 參考 Taiwan Llama Demo 網頁的預設參數
    data = {
        "inputs": prompt,
        "parameters": {
            "do_sample": True,
            "best_of": 1,
            "max_new_tokens": 1000,
            "stop": ["\n\n"],
            "temperature": 0.7,
            "top_k": 50,
            "top_p": 0.90,
        },
    }

    res = requests.post(url, json=data)
    return json.loads(res.text)[0]["generated_text"]


if __name__ == "__main__":
    main()

使用 RTX 3090 約四分半可以完成整份翻譯，給各位參考看看。

結論

今天介紹了 TGI 推論框架，是許多推論服務採用的框架，他的速度與廣泛模型支援、量化技術等，使其成為無論是開發者還是研究者的首選框架。

其實除了這幾天介紹的 ONNX, ggml, vLLM, TGI 以外，還有非常多的推論框架，例如 Exllama V2, CTranslate 2 等等，這些推論框架都有其穩定的開發與用戶。

不過這些框架的變化速度也相當快，很有可能筆者這幾天介紹的框架、參數和用法等等，隔天一個大改就全部不能用了。（好沒有技術保存價值的划水文章 Q_Q

接下來會介紹 Offloading 的概念，明天見！

參考

• GitHub: Text Generation Inference
• HF Docs: Text Generation Inference
• Text Generation Inference API
• GitHub: Taiwan Llama Issue#18